{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1aaab85c",
   "metadata": {},
   "source": [
    "# ``BrainPyObject`` and ``Collector``\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_advanced/base_and_collector.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_advanced/base_and_collector.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17e64f22",
   "metadata": {},
   "source": [
    "@[Chaoming Wang](https://github.com/chaoming0625)\n",
    "@[Xiaoyu Chen](mailto:c-xy17@tsinghua.org.cn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c07b0c",
   "metadata": {},
   "source": [
    "In this section, we are going to talk about:\n",
    "\n",
    "- The ``BrainPyObject`` class for the BrainPy ecosystem\n",
    "- The ``Collector`` to facilitate variable collection and manipulation."
   ]
  },
  {
   "cell_type": "code",
   "id": "1a9986eb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.417366Z",
     "start_time": "2025-10-06T03:28:31.911302Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "# bm.set_platform('cpu')"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "id": "45babeb2",
   "metadata": {},
   "source": [
    "## `brainpy.BrainPyObject`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa23d77e",
   "metadata": {},
   "source": [
    "The foundation of BrainPy is [brainpy.BrainPyObject](../apis/auto/generated/brainpy.base.BrainPyObject.rst). A BrainPyObject instance is an object which has variables and methods. All methods in the BrainPyObject object can be [JIT compiled](./compilation.ipynb) or [automatically differentiated](./differentiation.ipynb). In other words, any **class objects** that will be JIT compiled or automatically differentiated must inherent from ``brainpy.BrainPyObject``. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9e372bc",
   "metadata": {},
   "source": [
    "A BrainPyObject object can have many variables, children BrainPyObject objects, integrators, and methods. Below is the implemention of a [FitzHugh-Nagumo neuron model](https://brainmodels.readthedocs.io/en/brainpy-2.x/apis/generated/brainmodels.neurons.FHN.html) as an example. "
   ]
  },
  {
   "cell_type": "code",
   "id": "0c0b30c2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.429207Z",
     "start_time": "2025-10-06T03:28:35.423729Z"
    }
   },
   "source": [
    "class FHN(bp.BrainPyObject):\n",
    "  def __init__(self, num, a=0.7, b=0.8, tau=12.5, Vth=1.9, name=None):\n",
    "    super(FHN, self).__init__(name=name)\n",
    "\n",
    "    # parameters\n",
    "    self.num = num\n",
    "    self.a = a\n",
    "    self.b = b\n",
    "    self.tau = tau\n",
    "    self.Vth = Vth\n",
    "\n",
    "    # variables\n",
    "    self.V = bm.Variable(bm.zeros(num))\n",
    "    self.w = bm.Variable(bm.zeros(num))\n",
    "    self.spike = bm.Variable(bm.zeros(num, dtype=bool))\n",
    "\n",
    "    # integral\n",
    "    self.integral = bp.odeint(method='rk4', f=self.derivative)\n",
    "\n",
    "  def derivative(self, V, w, t, Iext):\n",
    "    dw = (V + self.a - self.b * w) / self.tau\n",
    "    dV = V - V * V * V / 3 - w + Iext\n",
    "    return dV, dw\n",
    "\n",
    "  def update(self, _t, _dt, x):\n",
    "    V, w = self.integral(self.V, self.w, _t, x)\n",
    "    self.spike[:] = bm.logical_and(V > self.Vth, self.V <= self.Vth)\n",
    "    self.w[:] = w\n",
    "    self.V[:] = V"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "id": "50e1d9f4",
   "metadata": {},
   "source": [
    "Note this model has three variables: ``self.V``, ``self.w``, and ``self.spike``. It also has an integrator ``self.integral``. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "530d0156",
   "metadata": {},
   "source": [
    "### The naming system"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de4f203d",
   "metadata": {},
   "source": [
    "Every BrainPyObject object has a unique name. Users can specify a unique name when you instantiate a BrainPyObject class. A used name will cause an error. "
   ]
  },
  {
   "cell_type": "code",
   "id": "e3c3be92",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.493943Z",
     "start_time": "2025-10-06T03:28:35.435397Z"
    }
   },
   "source": [
    "FHN(10, name='X').name"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'X'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "0d60e778",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.508402Z",
     "start_time": "2025-10-06T03:28:35.502565Z"
    }
   },
   "source": [
    "FHN(10, name='Y').name"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Y'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "3444aa16",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.539681Z",
     "start_time": "2025-10-06T03:28:35.531751Z"
    }
   },
   "source": [
    "try:\n",
    "    FHN(10, name='Y').name\n",
    "except Exception as e:\n",
    "    print(type(e).__name__, ':', e)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UniqueNameError : In BrainPy, each object should have a unique name. However, we detect that <__main__.FHN object at 0x0000022AAEBAD2E0> has a used name \"Y\". \n",
      "If you try to run multiple trials, you may need \n",
      "\n",
      ">>> brainpy.math.clear_name_cache() \n",
      "\n",
      "to clear all cached names. \n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "id": "af29b018",
   "metadata": {},
   "source": [
    "If a name is not specified to the BrainPyObject oject, BrainPy will assign a name for this object automatically. The rule for generating object name is ``class_name +  number_of_instances``. For example, ``FHN0``, ``FHN1``, etc."
   ]
  },
  {
   "cell_type": "code",
   "id": "9a4db12e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.576728Z",
     "start_time": "2025-10-06T03:28:35.570955Z"
    }
   },
   "source": [
    "FHN(10).name"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'FHN0'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "293c2a6e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.615844Z",
     "start_time": "2025-10-06T03:28:35.609541Z"
    }
   },
   "source": [
    "FHN(10).name"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'FHN1'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "id": "05ec44a4",
   "metadata": {},
   "source": [
    "Therefore, in BrainPy, you can access any object by its unique name, no matter how insignificant this object is."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "001de0c7",
   "metadata": {},
   "source": [
    "### Collection functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f64009e3",
   "metadata": {},
   "source": [
    "Three important collection functions are implemented for each BrainPyObject object. Specifically, they are:\n",
    "\n",
    "- ``nodes()``: to collect all instances of BrainPyObject objects, including children nodes in a node.\n",
    "- ``vars()``: to collect all variables defined in the BrainPyObject node and in its children nodes."
   ]
  },
  {
   "cell_type": "code",
   "id": "a21f85f5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.644270Z",
     "start_time": "2025-10-06T03:28:35.639887Z"
    }
   },
   "source": [
    "fhn = FHN(10)"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "id": "923cb0fc",
   "metadata": {},
   "source": [
    "All variables in a BrainPyObject object can be collected through ``BrainPyObject.vars()``. The returned container is a [ArrayCollector](../apis/auto/generated/brainpy.base.ArrayCollector.rst) (a subclass of ``Collector``)."
   ]
  },
  {
   "cell_type": "code",
   "id": "bc97484a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.673173Z",
     "start_time": "2025-10-06T03:28:35.668750Z"
    }
   },
   "source": [
    "vars = fhn.vars()\n",
    "\n",
    "vars"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'FHN2.V': Variable(\n",
       "   value=ShapedArray(float32[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN2.w': Variable(\n",
       "   value=ShapedArray(float32[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN2.spike': Variable(\n",
       "   value=ShapedArray(bool[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "id": "4996b36f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.701658Z",
     "start_time": "2025-10-06T03:28:35.698054Z"
    }
   },
   "source": [
    "type(vars)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "brainpy.math.object_transform.collectors.ArrayCollector"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "id": "e8dd1dbb",
   "metadata": {},
   "source": [
    "All nodes in the model can also be collected through one method ``BrainPyObject.nodes()``. The result container is an instance of [Collector](../apis/auto/generated/brainpy.base.Collector.rst)."
   ]
  },
  {
   "cell_type": "code",
   "id": "5a687cbf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.734177Z",
     "start_time": "2025-10-06T03:28:35.729724Z"
    }
   },
   "source": [
    "nodes = fhn.nodes()\n",
    "\n",
    "nodes  # note: integrator is also a node"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'FHN2': <__main__.FHN at 0x22aaebd6de0>,\n",
       " 'RK44': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aaebd76e0>}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "id": "e43569c6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.774154Z",
     "start_time": "2025-10-06T03:28:35.770086Z"
    }
   },
   "source": [
    "type(nodes)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "brainpy.math.object_transform.collectors.Collector"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "id": "7891c4e0",
   "metadata": {},
   "source": [
    "All integrators can be collected by:"
   ]
  },
  {
   "cell_type": "code",
   "id": "7f6ecf69",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.816730Z",
     "start_time": "2025-10-06T03:28:35.811860Z"
    }
   },
   "source": [
    "ints = fhn.nodes().subset(bp.integrators.Integrator)\n",
    "\n",
    "ints"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'RK44': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aaebd76e0>}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "id": "d489949b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.848695Z",
     "start_time": "2025-10-06T03:28:35.844151Z"
    }
   },
   "source": [
    "type(ints)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "brainpy.math.object_transform.collectors.Collector"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "id": "49eca479",
   "metadata": {},
   "source": [
    "Now, let's make a more complicated model by using the previously defined model ``FHN``. "
   ]
  },
  {
   "cell_type": "code",
   "id": "e636baa0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:35.876821Z",
     "start_time": "2025-10-06T03:28:35.872160Z"
    }
   },
   "source": [
    "class FeedForwardCircuit(bp.BrainPyObject):\n",
    "    def __init__(self, num1, num2, w=0.1, a=0.7, b=0.8, tau=12.5, Vth=1.9, name=None):\n",
    "        super(FeedForwardCircuit, self).__init__(name=name)\n",
    "        \n",
    "        self.pre = FHN(num1, a=a, b=b, tau=tau, Vth=Vth)\n",
    "        self.post = FHN(num2, a=a, b=b, tau=tau, Vth=Vth)\n",
    "        \n",
    "        self.conn = bm.ones((num1, num2), dtype=bool) * w\n",
    "        self.conn = bm.fill_diagonal(self.conn, 0., inplace=False)\n",
    "\n",
    "    def update(self, _t, _dt, x):\n",
    "        self.pre.update(_t, _dt, x)\n",
    "        x2 = self.pre.spike @ self.conn\n",
    "        self.post.update(_t, _dt, x2)"
   ],
   "outputs": [],
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "id": "2a0e9bcc",
   "metadata": {},
   "source": [
    "This model ``FeedForwardCircuit`` defines two layers. Each layer is modeled as a FitzHugh-Nagumo model (``FHN``). The first layer is densely connected to the second layer. The input to the second layer is the product of the first layer's spike and the connection strength ``w``. "
   ]
  },
  {
   "cell_type": "code",
   "id": "6aafd6c0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.174332Z",
     "start_time": "2025-10-06T03:28:35.905166Z"
    }
   },
   "source": [
    "net = FeedForwardCircuit(8, 5)"
   ],
   "outputs": [],
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "id": "54808da9",
   "metadata": {},
   "source": [
    "We can retrieve all variables by ``.vars()``:"
   ]
  },
  {
   "cell_type": "code",
   "id": "c2f7692a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.186634Z",
     "start_time": "2025-10-06T03:28:36.182038Z"
    }
   },
   "source": [
    "net.vars()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'FHN3.V': Variable(\n",
       "   value=ShapedArray(float32[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN3.w': Variable(\n",
       "   value=ShapedArray(float32[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN3.spike': Variable(\n",
       "   value=ShapedArray(bool[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN4.V': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN4.w': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN4.spike': Variable(\n",
       "   value=ShapedArray(bool[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 17
  },
  {
   "cell_type": "markdown",
   "id": "4c9eec24",
   "metadata": {},
   "source": [
    "And retrieve all nodes (instances of the BrainPyObject class) by ``.nodes()``:"
   ]
  },
  {
   "cell_type": "code",
   "id": "bd98a238",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.218743Z",
     "start_time": "2025-10-06T03:28:36.214336Z"
    }
   },
   "source": [
    "net.nodes()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'FeedForwardCircuit0': <__main__.FeedForwardCircuit at 0x22aae606030>,\n",
       " 'FHN3': <__main__.FHN at 0x22aaebd6060>,\n",
       " 'FHN4': <__main__.FHN at 0x22aaebd5f10>,\n",
       " 'RK45': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe61550>,\n",
       " 'RK46': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe63410>}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 18
  },
  {
   "cell_type": "markdown",
   "id": "8d859dca",
   "metadata": {},
   "source": [
    "If we only care about a subtype of class, we can retrieve them through:"
   ]
  },
  {
   "cell_type": "code",
   "id": "89dc3aeb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.275461Z",
     "start_time": "2025-10-06T03:28:36.270611Z"
    }
   },
   "source": [
    "net.nodes().subset(bp.ode.ODEIntegrator)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'RK45': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe61550>,\n",
       " 'RK46': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe63410>}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 19
  },
  {
   "cell_type": "markdown",
   "id": "32906d25",
   "metadata": {},
   "source": [
    "#### Absolute paths\n",
    "\n",
    "It's worthy to note that there are two ways to access variables, integrators, and nodes. They are \"absolute\" paths and \"relative\" paths. The default way is the absolute path. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1abda9ea",
   "metadata": {},
   "source": [
    "For absolute paths, all keys in the resulting Collector (``BrainPyObject.nodes()``) has the format of ``key = node_name [+ field_name]``. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce8ea30",
   "metadata": {},
   "source": [
    "**.nodes() example 1**: In the above ``fhn`` instance, there are two nodes: \"fnh\" and its integrator \"fhn.integral\"."
   ]
  },
  {
   "cell_type": "code",
   "id": "1b947d76",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.313799Z",
     "start_time": "2025-10-06T03:28:36.310144Z"
    }
   },
   "source": [
    "fhn.integral.name, fhn.name"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('RK44', 'FHN2')"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 20
  },
  {
   "cell_type": "markdown",
   "id": "434edabf",
   "metadata": {},
   "source": [
    "Calling ``.nodes()`` returns their names and models. "
   ]
  },
  {
   "cell_type": "code",
   "id": "05faecf5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.349664Z",
     "start_time": "2025-10-06T03:28:36.345360Z"
    }
   },
   "source": [
    "fhn.nodes().keys()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['FHN2', 'RK44'])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 21
  },
  {
   "cell_type": "markdown",
   "id": "cc3fb400",
   "metadata": {},
   "source": [
    "**.nodes() example 2**: In the above ``net`` instance, there are five nodes:"
   ]
  },
  {
   "cell_type": "code",
   "id": "bc3b0041",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.382308Z",
     "start_time": "2025-10-06T03:28:36.377535Z"
    }
   },
   "source": [
    "net.pre.name, net.post.name, net.pre.integral.name, net.post.integral.name, net.name"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('FHN3', 'FHN4', 'RK45', 'RK46', 'FeedForwardCircuit0')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 22
  },
  {
   "cell_type": "markdown",
   "id": "37e3835d",
   "metadata": {},
   "source": [
    "Calling ``.nodes()`` also returns the names and instances of all models. "
   ]
  },
  {
   "cell_type": "code",
   "id": "a6757d5e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.410591Z",
     "start_time": "2025-10-06T03:28:36.406053Z"
    }
   },
   "source": [
    "net.nodes().keys()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['FeedForwardCircuit0', 'FHN3', 'FHN4', 'RK45', 'RK46'])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 23
  },
  {
   "cell_type": "markdown",
   "id": "fe07b062",
   "metadata": {},
   "source": [
    "**.vars() example 1**: In the above ``fhn`` instance, there are three variables: \"V\", \"w\" and \"input\". Calling ``.vars()`` returns a dict of `<node_name + var_name, var_value>`. "
   ]
  },
  {
   "cell_type": "code",
   "id": "e7b0014e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.442309Z",
     "start_time": "2025-10-06T03:28:36.438379Z"
    }
   },
   "source": [
    "fhn.vars()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'FHN2.V': Variable(\n",
       "   value=ShapedArray(float32[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN2.w': Variable(\n",
       "   value=ShapedArray(float32[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN2.spike': Variable(\n",
       "   value=ShapedArray(bool[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 24
  },
  {
   "cell_type": "markdown",
   "id": "146bf744",
   "metadata": {},
   "source": [
    "**.vars() example 2**: This also applies in the ``net`` instance:"
   ]
  },
  {
   "cell_type": "code",
   "id": "1696e602",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.474974Z",
     "start_time": "2025-10-06T03:28:36.469980Z"
    }
   },
   "source": [
    "net.vars()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'FHN3.V': Variable(\n",
       "   value=ShapedArray(float32[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN3.w': Variable(\n",
       "   value=ShapedArray(float32[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN3.spike': Variable(\n",
       "   value=ShapedArray(bool[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN4.V': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN4.w': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'FHN4.spike': Variable(\n",
       "   value=ShapedArray(bool[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 25
  },
  {
   "cell_type": "markdown",
   "id": "8381a0b0",
   "metadata": {},
   "source": [
    "#### Relative paths"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79dbb038",
   "metadata": {},
   "source": [
    "Variables, integrators, and nodes can also be accessed by relative paths. For example, the ``pre`` instance in the ``net`` can be accessed by"
   ]
  },
  {
   "cell_type": "code",
   "id": "8d921a96",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.522572Z",
     "start_time": "2025-10-06T03:28:36.518620Z"
    }
   },
   "source": [
    "net.pre"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.FHN at 0x22aaebd6060>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 26
  },
  {
   "cell_type": "markdown",
   "id": "c041cd24",
   "metadata": {},
   "source": [
    "Relative paths preserve the dependence relationship. For example, all nodes retrieved from the perspective of ``net`` are:"
   ]
  },
  {
   "cell_type": "code",
   "id": "c8c1cb0e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.569453Z",
     "start_time": "2025-10-06T03:28:36.564527Z"
    }
   },
   "source": [
    "net.nodes(method='relative')"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'': <__main__.FeedForwardCircuit at 0x22aae606030>,\n",
       " 'pre': <__main__.FHN at 0x22aaebd6060>,\n",
       " 'post': <__main__.FHN at 0x22aaebd5f10>,\n",
       " 'pre.integral': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe61550>,\n",
       " 'post.integral': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe63410>}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 27
  },
  {
   "cell_type": "markdown",
   "id": "60fd0239",
   "metadata": {},
   "source": [
    "However, nodes retrieved from the start point of ``net.pre`` will be:"
   ]
  },
  {
   "cell_type": "code",
   "id": "ac50b5b5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.602673Z",
     "start_time": "2025-10-06T03:28:36.598930Z"
    }
   },
   "source": [
    "net.pre.nodes('relative')"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'': <__main__.FHN at 0x22aaebd6060>,\n",
       " 'integral': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafe61550>}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 28
  },
  {
   "cell_type": "markdown",
   "id": "ca6b200f",
   "metadata": {},
   "source": [
    "Variables can also br relatively inferred from the model. For example, variables that can be relatively accessed from ``net`` include:"
   ]
  },
  {
   "cell_type": "code",
   "id": "7ec908d2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.630605Z",
     "start_time": "2025-10-06T03:28:36.625214Z"
    }
   },
   "source": [
    "net.vars('relative')"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'pre.V': Variable(\n",
       "   value=ShapedArray(float32[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'pre.w': Variable(\n",
       "   value=ShapedArray(float32[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'pre.spike': Variable(\n",
       "   value=ShapedArray(bool[8]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'post.V': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'post.w': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'post.spike': Variable(\n",
       "   value=ShapedArray(bool[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "cell_type": "markdown",
   "id": "16fa7a2f",
   "metadata": {},
   "source": [
    "While variables relatively accessed from ``net.post`` are:"
   ]
  },
  {
   "cell_type": "code",
   "id": "1f77f979",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.661967Z",
     "start_time": "2025-10-06T03:28:36.657712Z"
    }
   },
   "source": [
    "net.post.vars('relative')"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'V': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'w': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'spike': Variable(\n",
       "   value=ShapedArray(bool[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 30
  },
  {
   "cell_type": "markdown",
   "id": "e63c7c80",
   "metadata": {},
   "source": [
    "####  Elements in containers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8aec9ca",
   "metadata": {},
   "source": [
    "One drawback of collection functions is that they don not look for elements in *list*, *dict* or any other container structure. "
   ]
  },
  {
   "cell_type": "code",
   "id": "4544d9fe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.689511Z",
     "start_time": "2025-10-06T03:28:36.684491Z"
    }
   },
   "source": [
    "class ATest(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(ATest, self).__init__()\n",
    "        \n",
    "        self.all_vars = [bm.Variable(bm.zeros(5)), bm.Variable(bm.ones(6)),]\n",
    "        self.sub_nodes = {'a': FHN(10), 'b': FHN(5)}"
   ],
   "outputs": [],
   "execution_count": 31
  },
  {
   "cell_type": "code",
   "id": "b1199b05",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.730592Z",
     "start_time": "2025-10-06T03:28:36.711431Z"
    }
   },
   "source": [
    "t1 = ATest()"
   ],
   "outputs": [],
   "execution_count": 32
  },
  {
   "cell_type": "markdown",
   "id": "3d3bcbac",
   "metadata": {},
   "source": [
    "The above class defines a list of variables, and a dict of children nodes, but the variables and children nodes cannot be retrieved from the collection functions ``vars()`` and ``nodes()``. "
   ]
  },
  {
   "cell_type": "code",
   "id": "c4efd7c6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.741781Z",
     "start_time": "2025-10-06T03:28:36.737605Z"
    }
   },
   "source": [
    "t1.vars()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 33
  },
  {
   "cell_type": "code",
   "id": "df9fd51b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.773541Z",
     "start_time": "2025-10-06T03:28:36.769422Z"
    }
   },
   "source": [
    "t1.nodes()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ATest0': <__main__.ATest at 0x22aafe61af0>}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 34
  },
  {
   "cell_type": "markdown",
   "id": "23f81f1e",
   "metadata": {},
   "source": [
    "To solve this problem, BrianPy provides ``implicit_vars`` and ``implicit_nodes`` (an instance of \"dict\") to hold variables and nodes in container structures. Variables registered in ``implicit_vars`` and integrators and nodes registered in ``implicit_nodes`` can be retrieved by collection functions."
   ]
  },
  {
   "cell_type": "code",
   "id": "767f2f74",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.799790Z",
     "start_time": "2025-10-06T03:28:36.793806Z"
    }
   },
   "source": [
    "class AnotherTest(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(AnotherTest, self).__init__()\n",
    "        \n",
    "        self.all_vars = [bm.Variable(bm.zeros(5)), bm.Variable(bm.ones(6)),]\n",
    "        self.sub_nodes = {'a': FHN(10, name='T1'), 'b': FHN(5, name='T2')}\n",
    "        \n",
    "        self.register_implicit_vars({f'v{i}': v for i, v in enumerate(self.all_vars)})  # the input must be a dict\n",
    "        self.register_implicit_nodes({k: v for k, v in self.sub_nodes.items()})  # the input must be a dict"
   ],
   "outputs": [],
   "execution_count": 35
  },
  {
   "cell_type": "code",
   "id": "e5fc3ec0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.836337Z",
     "start_time": "2025-10-06T03:28:36.830221Z"
    }
   },
   "source": [
    "t2 = AnotherTest()"
   ],
   "outputs": [],
   "execution_count": 36
  },
  {
   "cell_type": "code",
   "id": "55ee8df8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.858611Z",
     "start_time": "2025-10-06T03:28:36.854593Z"
    }
   },
   "source": [
    "# This model has two \"FHN\" instances, each \"FHN\" instance has one integrator. \n",
    "# Therefore, there are five BrainPyObject objects. \n",
    "\n",
    "t2.nodes()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'AnotherTest0': <__main__.AnotherTest at 0x22aae62d700>,\n",
       " 'T1': <__main__.FHN at 0x22aaebd7bc0>,\n",
       " 'T2': <__main__.FHN at 0x22aafef60f0>,\n",
       " 'RK49': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafef6150>,\n",
       " 'RK410': <brainpy.integrators.ode.explicit_rk.RK4 at 0x22aafef6510>}"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 37
  },
  {
   "cell_type": "code",
   "id": "3740e867",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:28:36.903921Z",
     "start_time": "2025-10-06T03:28:36.898765Z"
    }
   },
   "source": [
    "# This model has two FHN node, each of which has three variables.\n",
    "# Moreover, this model has two implicit variables.\n",
    "\n",
    "t2.vars()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'AnotherTest0.v0': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'AnotherTest0.v1': Variable(\n",
       "   value=ShapedArray(float32[6]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'T1.V': Variable(\n",
       "   value=ShapedArray(float32[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'T1.w': Variable(\n",
       "   value=ShapedArray(float32[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'T1.spike': Variable(\n",
       "   value=ShapedArray(bool[10]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'T2.V': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'T2.w': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'T2.spike': Variable(\n",
       "   value=ShapedArray(bool[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 38
  },
  {
   "cell_type": "markdown",
   "id": "d38d5954",
   "metadata": {},
   "source": [
    "## Collector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886dab02",
   "metadata": {},
   "source": [
    "Collection functions return an ``brainpy.Collector`` that is a dictionary mapping names to elements. It has some useful methods. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b61c9ced",
   "metadata": {},
   "source": [
    "### ``subset()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30d0ff65",
   "metadata": {},
   "source": [
    "``Collector.subset(cls)`` returns a part of elements whose type is the given ``cls``. For example, ``BrainPyObject.nodes()`` returns all instances of BrainPyObject class. If you are only interested in one type, like ``ODEIntegrator``, you can use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8d3f0dea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'RK45': <brainpy.integrators.ode.explicit_rk.RK4 at 0x13bea139a00>,\n",
       " 'RK46': <brainpy.integrators.ode.explicit_rk.RK4 at 0x13bea13bcd0>}"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.nodes().subset(bp.ode.ODEIntegrator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "371f4933",
   "metadata": {},
   "source": [
    "Actually, ``Collector.subset(cls)`` travels all the elements in this collection, and find the element whose type matches the given ``cls``. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e4226a",
   "metadata": {},
   "source": [
    "### ``unique()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0d3e4e7",
   "metadata": {},
   "source": [
    "It is common in machine learning that weights are shared with several objects, or the same weight can be accessed by various dependence relationships. Collection functions of BrainPyObject usually return a collection in which the same value have multiple keys. The duplicate elements will not be automatically excluded. However, it is important not to apply operations such as gradient descent twice or more to the same elements. \n",
    "\n",
    "Therefore, the Collector provides ``Collector.unique()`` to handle this problem automatically. ``Collector.unique()`` returns a copy of collection in which all elements are unique. "
   ]
  },
  {
   "cell_type": "code",
   "id": "73870138",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:33.862531Z",
     "start_time": "2025-10-06T03:29:33.856271Z"
    }
   },
   "source": [
    "class ModelA(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(ModelA, self).__init__()\n",
    "        self.a = bm.Variable(bm.zeros(5))\n",
    "\n",
    "        \n",
    "class SharedA(bp.BrainPyObject):\n",
    "    def __init__(self, source):\n",
    "        super(SharedA, self).__init__()\n",
    "        self.source = source\n",
    "        self.a = source.a  # shared variable\n",
    "        \n",
    "        \n",
    "class Group(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(Group, self).__init__()\n",
    "        self.A = ModelA()\n",
    "        self.A_shared = SharedA(self.A)\n",
    "\n",
    "g = Group()"
   ],
   "outputs": [],
   "execution_count": 40
  },
  {
   "cell_type": "code",
   "id": "0b975732",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:34.776154Z",
     "start_time": "2025-10-06T03:29:34.770104Z"
    }
   },
   "source": [
    "g.vars('relative')  # save Variable can be accessed by three paths"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'A.a': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'A_shared.a': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " ),\n",
       " 'A_shared.source.a': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 41
  },
  {
   "cell_type": "code",
   "id": "fc55069f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:35.606721Z",
     "start_time": "2025-10-06T03:29:35.602049Z"
    }
   },
   "source": [
    "g.vars('relative').unique()  # only return a unique path"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'A.a': Variable(\n",
       "   value=ShapedArray(float32[5]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 42
  },
  {
   "cell_type": "code",
   "id": "898f898b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:36.256335Z",
     "start_time": "2025-10-06T03:29:36.252164Z"
    }
   },
   "source": [
    "g.nodes('relative')  # \"ModelA\" is accessed twice"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'': <__main__.Group at 0x22ab12c22a0>,\n",
       " 'A': <__main__.ModelA at 0x22ab12c3e00>,\n",
       " 'A_shared': <__main__.SharedA at 0x22ab12c6b10>,\n",
       " 'A_shared.source': <__main__.ModelA at 0x22ab12c3e00>}"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 43
  },
  {
   "cell_type": "code",
   "id": "532bd337",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:36.693727Z",
     "start_time": "2025-10-06T03:29:36.689412Z"
    }
   },
   "source": [
    "g.nodes('relative').unique()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'': <__main__.Group at 0x22ab12c22a0>,\n",
       " 'A': <__main__.ModelA at 0x22ab12c3e00>,\n",
       " 'A_shared': <__main__.SharedA at 0x22ab12c6b10>}"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 44
  },
  {
   "cell_type": "markdown",
   "id": "942093ce",
   "metadata": {},
   "source": [
    "### ``update()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23f70fce",
   "metadata": {},
   "source": [
    "The Collector can also catch potential conflicts during the assignment. The bracket assignment of a Collector (``[key]``) and ``Collector.update()`` will check whether the same key is mapped to a different value. If it is, an error will occur. "
   ]
  },
  {
   "cell_type": "code",
   "id": "9cf6c7f6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:38.131427Z",
     "start_time": "2025-10-06T03:29:38.126811Z"
    }
   },
   "source": [
    "tc = bp.Collector({'a': bm.zeros(10)})\n",
    "\n",
    "tc"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': Array(value=Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), dtype=float32)}"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 45
  },
  {
   "cell_type": "code",
   "id": "db8804d7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:38.883438Z",
     "start_time": "2025-10-06T03:29:38.866286Z"
    }
   },
   "source": [
    "try:\n",
    "    tc['a'] = bm.zeros(1)  # same key \"a\", different tensor\n",
    "except Exception as e:\n",
    "    print(type(e).__name__, \":\", e)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError : Name \"a\" conflicts: same name for [0.] and [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.].\n"
     ]
    }
   ],
   "execution_count": 46
  },
  {
   "cell_type": "code",
   "id": "531922ec",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:39.317657Z",
     "start_time": "2025-10-06T03:29:39.313418Z"
    }
   },
   "source": [
    "try:\n",
    "    tc.update({'a': bm.ones(1)})  # same key \"a\", different tensor\n",
    "except Exception as e:\n",
    "    print(type(e).__name__, \":\", e)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError : Name \"a\" conflicts: same name for [1.] and [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.].\n"
     ]
    }
   ],
   "execution_count": 47
  },
  {
   "cell_type": "markdown",
   "id": "15fc1991",
   "metadata": {},
   "source": [
    "### ``replace()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bb2976c",
   "metadata": {},
   "source": [
    "``Collector.replace(old_key, new_value)`` is used to update the value of a key. "
   ]
  },
  {
   "cell_type": "code",
   "id": "c91f4373",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:40.771110Z",
     "start_time": "2025-10-06T03:29:40.766270Z"
    }
   },
   "source": [
    "tc"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': Array(value=Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), dtype=float32)}"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 48
  },
  {
   "cell_type": "code",
   "id": "ecf2a15c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:41.435464Z",
     "start_time": "2025-10-06T03:29:41.413500Z"
    }
   },
   "source": [
    "tc.replace('a', bm.ones(3))\n",
    "\n",
    "tc"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': Array(value=Array([1., 1., 1.]), dtype=float32)}"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 49
  },
  {
   "cell_type": "markdown",
   "id": "6401d334",
   "metadata": {},
   "source": [
    "### ``__add()__``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c2b83ee",
   "metadata": {},
   "source": [
    "Two Collectors can be merged. "
   ]
  },
  {
   "cell_type": "code",
   "id": "93bde37c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:29:42.900790Z",
     "start_time": "2025-10-06T03:29:42.894898Z"
    }
   },
   "source": [
    "a = bp.Collector({'a': bm.zeros(10)})\n",
    "b = bp.Collector({'b': bm.ones(10)})\n",
    "\n",
    "a + b"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': Array(value=Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), dtype=float32),\n",
       " 'b': Array(value=Array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]), dtype=float32)}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 50
  },
  {
   "cell_type": "markdown",
   "id": "7cb93e3d",
   "metadata": {},
   "source": [
    "## ArrayCollector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cbd07f6",
   "metadata": {},
   "source": [
    "``ArrayCollector`` is subclass of ``Collector``, but it is specifically to collect arrays."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "279.273px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
